package com.whosly.rapid.lang.thread.executor;

import java.lang.reflect.Field;
import java.util.Iterator;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executors;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import com.whosly.rapid.lang.thread.AsyncLoadCallable;
import com.whosly.rapid.lang.thread.AsyncLoadFuture;
import com.whosly.rapid.lang.thread.factory.INamedThreadFactory;
import com.whosly.rapid.lang.thread.factory.NamedThreadFactory;
import com.whosly.rapid.lang.thread.hook.ShutdownHookService;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.StringUtils;

import com.google.common.collect.MapMaker;
import org.springframework.util.ReflectionUtils;

/**
 * 线程池，增加线程池的监控，定时的扫描线程池状态<br>
 * 该实例会主动注册 ThreadPoolRegistry.getInstance().register<br>
 * 对应 MonitorThreadPoolExecutor  <br>
 *
 * 扩展了J.U.C的ThreadPoolExecutor，主要扩展点说明：
 *
 * <pre>
 * 1. 覆写newTaskFor函数，返回自定义的{@linkplain AsyncLoadFuture}
 * 2. 增强了Pool池中的Worker线程，会自动复制caller Thread的threadLocal信息，几点考虑：
 *   a. Worker线程为pool的内部管理对象，在操作ThreadLocal信息时安全性上不存在问题，持有的引用在task完成后也可以正常释放。ThreadLocal引用在Worker线程中的生命周期<=Caller Thread线程
 *   b. 做为并行异步加载，一个主要的设计思想就是对业务尽可能的透明，尽可能的减少使用陷井，所以这里通过非正常手段实现了ThreadLocal的支持，实属无奈
 * </pre>
 *
 */
public class MonitorThreadPoolExecutor extends ThreadPoolExecutor {
    private static final Field THREAD_LOCAL_FIELD = ReflectionUtils.findField(Thread.class, "threadLocals");
    private static final Field INHERITABLE_THREAD_LOCAL_FIELD = ReflectionUtils.findField(Thread.class,
            "inheritableThreadLocals");

    static {
        // 强制的声明accessible
        ReflectionUtils.makeAccessible(THREAD_LOCAL_FIELD);
        ReflectionUtils.makeAccessible(INHERITABLE_THREAD_LOCAL_FIELD);
    }

    // ======================= 继承自ThreadPoolExecutor的构造函数 =======================
    /**
     * 该实例会主动注册 ShutdownHookService.register<br>
     * <strong>核心线程池大小和最大线程池大小均为1 <br>
     * 最大空闲时间默认为0 </strong> 
     * 
     * @param workQueue 任务队列
     */
    public MonitorThreadPoolExecutor(String namePrefix, BlockingQueue<Runnable> workQueue) {
    	this(namePrefix, 0L, TimeUnit.MILLISECONDS, workQueue);
    }
    
    /**
     * 该实例会主动注册 ShutdownHookService.register<br>
     * <strong> 核心线程池大小和最大线程池大小均为 1 </strong><br>
     * 
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     */
    public MonitorThreadPoolExecutor(String namePrefix, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
    	this(namePrefix, 1, 1, keepAliveTime, unit, workQueue);
    }

    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     * @param handler 拒绝策略。 默认  new AbortPolicy()
     */
    public MonitorThreadPoolExecutor(String namePrefix, int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
    	this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, new NamedThreadFactory(namePrefix),  handler);
    }
    
    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     * @param handler 拒绝策略。 默认  new AbortPolicy()
     */
    public MonitorThreadPoolExecutor(String namePrefix, int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler, boolean daemon) {
    	this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, new NamedThreadFactory(namePrefix, daemon),  handler);
    }
    
    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     * @param threadFactory  线程创建工厂
     * @param handler 拒绝策略
     */
    public MonitorThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, INamedThreadFactory threadFactory,
                                 RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
                threadFactory, handler);
        ThreadPoolRegistry.getInstance().registerThreadPool(this);
    }

    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     * @param threadFactory 线程创建工厂
     */
    public MonitorThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, INamedThreadFactory threadFactory) {
        this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
                threadFactory, new RejectedPolicyHandler(threadFactory.getNamePrefix()));
    }
    
    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     */
    public MonitorThreadPoolExecutor(String namePrefix, int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
    	this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, new NamedThreadFactory(namePrefix), new RejectedPolicyHandler(namePrefix));
    }
    
    /**
     * 该实例会主动注册 ShutdownHookService.register
     * 
     * @param corePoolSize 核心线程池大小
     * @param maximumPoolSize 最大线程池大小
     * @param keepAliveTime 超出核心线程池大小的线程最大空闲时间
     * @param unit 时间单位
     * @param workQueue 任务队列
     */
    public MonitorThreadPoolExecutor(String namePrefix, int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue,  boolean daemon) {
    	this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, new NamedThreadFactory(namePrefix, daemon));
    }


    // ======================= 扩展点 =======================
    @Override
    protected void beforeExecute(Thread t, Runnable command) {
        // 在执行之前处理下ThreadPool的属性继承
        if (command instanceof AsyncLoadFuture) {
            AsyncLoadFuture<?> afuture = (AsyncLoadFuture<?>) command;
            boolean flag = afuture.getConfig().getNeedThreadLocalSupport();
            if (flag) {
                initThreadLocal(THREAD_LOCAL_FIELD, afuture.getCallerThread(), t);
                initThreadLocal(INHERITABLE_THREAD_LOCAL_FIELD, afuture.getCallerThread(), t);
            }
        }

        super.beforeExecute(t, command);
    }

    @Override
    protected void afterExecute(Runnable command, Throwable t) {
        // 在执行结束后清理下ThreadPool的属性，GC处理
        if (command instanceof AsyncLoadFuture) {
            AsyncLoadFuture<?> afuture = (AsyncLoadFuture<?>) command;
            boolean flag = afuture.getConfig().getNeedThreadLocalSupport();
            if (flag) {
                recoverThreadLocal(THREAD_LOCAL_FIELD, afuture.getCallerThread(), afuture.getRunnerThread());
                recoverThreadLocal(INHERITABLE_THREAD_LOCAL_FIELD, afuture.getCallerThread(), afuture.getRunnerThread());
            }
        }

        super.afterExecute(command, t);
    }

    public <T> AsyncLoadFuture<T> submit(AsyncLoadCallable<T> task) {
        if (task == null) {
            throw new NullPointerException();
        }

        // 使用自定义的Future
        AsyncLoadFuture<T> ftask = new AsyncLoadFuture<T>(task);
        execute(ftask);
        return ftask;
    }

    @Override
    public void execute(Runnable command) {
        if (command instanceof AsyncLoadFuture) {
            AsyncLoadFuture<?> afuture = (AsyncLoadFuture<?>) command;
            boolean flag = afuture.getConfig().getNeedThreadLocalSupport();
            if (flag) {
                Thread thread = Thread.currentThread();
                if (ReflectionUtils.getField(THREAD_LOCAL_FIELD, thread) == null) {
                    // 创建一个空的ThreadLocal,立马写回去
                    new ThreadLocal<Boolean>(); // 这时会在runner线程产生一空记录的ThreadLocalMap记录
                }
                if (ReflectionUtils.getField(INHERITABLE_THREAD_LOCAL_FIELD, thread) == null) {
                    // 创建一个空的ThreadLocal,立马写回去
                    new InheritableThreadLocal<Boolean>(); // 可继承的ThreadLocal
                }
            }
        }

        super.execute(command);// 调用父类进行提交
    }

    private void initThreadLocal(Field field, Thread caller, Thread runner) {
        if (caller == null || runner == null) {
            return;
        }
        // 主要考虑这样的情况：
        // 1.
        // 如果caller线程没有使用ThreadLocal对象，而异步加载的runner线程执行中使用了ThreadLocal对象，则需要复制对象到caller线程上
        // 2.
        // 后续caller,多个runner线程有使用ThreadLocal对象，使用的是同一个引用,直接set都是针对同一个ThreadLocal,所以以后就不需要进行合并

        // 因为在提交Runnable时已经同步创建了一个ThreadLocalMap对象，所以runner线程只需要复制caller对应的引用即可，不需要进行合并，简化处理
        // threadlocal属性复制,注意是引用复制
        Object callerThreadLocalMap = ReflectionUtils.getField(field, caller);
        if (callerThreadLocalMap != null) {
            ReflectionUtils.setField(field, runner, callerThreadLocalMap);// 复制caller的信息到runner线程上
        } else {
            // 这个分支不会出现,因为在execute提交的时候已经添加
        }
    }

    private void recoverThreadLocal(Field field, Thread caller, Thread runner) {
        if (runner == null) {
            return;
        }
        // 清理runner线程的ThreadLocal，为下一个task服务
        ReflectionUtils.setField(field, runner, null);
    }

    // ======================= 内部类 =======================
    /**
     * 线程池注册类
     */
    @Slf4j
    static class ThreadPoolRegistry {
    	/**
         * 默认监听线程检查周期
         */
        private static final Long DEFAULT_PERIOD = 60L;
        /**
         * 默认监听线程前缀
         */
        private static final String DEFAULT_NAME_PREFIX = "thread-pool-monitor-thread";
        
        // 使用WeakValueMap，当Value被垃圾回收时会将此value在map中的entry清除，防止内存溢出
        private static ConcurrentMap<String, ThreadPoolExecutor> threadPoolExecutorConcurrentMap =
                new MapMaker().weakValues().makeMap();


        private static ThreadPoolRegistry INSTANCE = new ThreadPoolRegistry();

        private AtomicLong atomicLong = new AtomicLong(0L);
        
        private ThreadPoolRegistry() {
            ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(new NamedThreadFactory(DEFAULT_NAME_PREFIX));
            scheduledExecutorService.scheduleAtFixedRate(new ThreadPoolMonitorTask(), 0L, DEFAULT_PERIOD, TimeUnit.SECONDS);

            ShutdownHookService.register(scheduledExecutorService);// 注册关闭钩子
        }

        public static ThreadPoolRegistry getInstance() {
            return INSTANCE;
        }

        public void registerThreadPool(ThreadPoolExecutor threadPoolExecutor) {

            threadPoolExecutorConcurrentMap.putIfAbsent(generateThreadPoolName(threadPoolExecutor), threadPoolExecutor);

            ShutdownHookService.register(threadPoolExecutor);
        }

        private String generateThreadPoolName(ThreadPoolExecutor threadPoolExecutor) {
            String threadPoolNamePrefix = null;
            if (threadPoolExecutor.getThreadFactory() instanceof INamedThreadFactory) {
                threadPoolNamePrefix = ((INamedThreadFactory) threadPoolExecutor.getThreadFactory()).getNamePrefix();
            }

            threadPoolNamePrefix = StringUtils.isNotBlank(threadPoolNamePrefix) ? "thread-pool-" + threadPoolNamePrefix + "-" : "thread-pool-default-";
            return String.format(threadPoolNamePrefix + "%d", atomicLong.incrementAndGet());
        }

        class ThreadPoolMonitorTask implements Runnable {

            @Override
            public void run() {
                try{
                    Iterator<String> iterator = threadPoolExecutorConcurrentMap.keySet().iterator();
                    while (iterator.hasNext()) {
                        String threadPoolName = iterator.next();
                        ThreadPoolExecutor threadPoolExecutor = threadPoolExecutorConcurrentMap.get(threadPoolName);
                        if (threadPoolExecutor.isTerminated()) {
                            iterator.remove();
                            log.warn("thread-pool {} is terminated", threadPoolName);
                            continue;
                        }
                        threadPoolSnapShot(threadPoolName, threadPoolExecutor);
                    }
                }catch (Exception ex){
                    log.error("thread pool monitor error", ex);
                }

            }

            /**
             * 数据搜集<br>
             * 监控指标: 
             * <pre>
             * thread pool name: 线程池名称，根据MizThreadFactory的前缀加上递增数字<code>AtomicLong</code>生成
             * poolSize: 线程池的线程数量，只要不销毁，只增不减
             * corePoolSize: 核心线程池大小
             * activeCount: 活动线程数量
             * taskCount: 线程池需要执行的任务总数量
             * completedTaskCount: 已经完成的任务数量
             * isShutDown: 是否被关闭
             * isTerminated: 线程池是否终止
             * 
             * 未添加的：
             * largestPoolSize:线程池曾经创建过的最大线程数量，可以推测出是否达到corePoolSize和MaxinumPoolSize
             * 当前排队线程数
             * </pre>
             */
            private void threadPoolSnapShot(String threadPoolName, ThreadPoolExecutor threadPoolExecutor) {
                log.info("thread pool name {}, poolSize {}, corePoolSize {}, activeCount {}, taskCount {}, completedTaskCount {}, isShutDown {}, isTerminated {}",
                        threadPoolName, threadPoolExecutor.getPoolSize(), threadPoolExecutor.getCorePoolSize(),
                        threadPoolExecutor.getActiveCount(), threadPoolExecutor.getTaskCount(), threadPoolExecutor.getCompletedTaskCount(),
                        threadPoolExecutor.isShutdown(), threadPoolExecutor.isTerminated());
            }
        }
    }
    
    /**
     * 线程池自定义拒绝策略，不对外访问
     * 
     * @author yueny09 <yueny09@163.com>
     *
     * @DATE 2018年9月20日 下午6:50:50
     */
    @Slf4j
    static class RejectedPolicyHandler extends AbortPolicy implements RejectedExecutionHandler {
    	private static AtomicLong totals = new AtomicLong(0L);
        
    	/**
    	 * 业务名，可以为空
    	 */
    	@Getter
        private String bizName;
          
        public RejectedPolicyHandler(String bizName){
        	this.bizName = bizName;
        }
          
        @Override  
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
        	//业务报警
        	StringBuilder sb = new StringBuilder("[");
        	sb.append(bizName);
        	sb.append("] 线程忙，请求被拒绝.max: ");
        	sb.append(executor.getMaximumPoolSize());
        	sb.append(", queue: ");
        	sb.append(executor.getQueue().size());
        	
            String warnText = sb.toString();  
            log.warn(warnText);  
            
            totals.addAndGet(1);  
            
            super.rejectedExecution(r, executor);  
        }  
      
        public long getRejectedSize() {  
            return totals.get();  
        }
        
    }

}
